#pragma once

#include <string>
#include <vector>
#include <memory>
#include <variant>

namespace MiniDB
{

    // 前向声明
    class Statement;
    class Expression;

    // 基类
    class ASTNode
    {
    public:
        virtual ~ASTNode() = default;
        virtual std::string toString(int indent = 0) const = 0;
    };

    // 语句节点
    class Statement : public ASTNode
    {
    };

    // 表达式节点
    class Expression : public ASTNode
    {
    };

    // 具体语句类型
    class SelectStatement : public Statement
    {
    public:
        std::vector<std::shared_ptr<Expression>> select_list;
        std::shared_ptr<ASTNode> from_clause;
        std::shared_ptr<Expression> where_clause;

        SelectStatement(std::vector<std::shared_ptr<Expression>> select_list_,
                        std::shared_ptr<ASTNode> from_clause_,
                        std::shared_ptr<Expression> where_clause_ = nullptr)
            : select_list(std::move(select_list_)),
              from_clause(std::move(from_clause_)),
              where_clause(std::move(where_clause_)) {}

        std::string toString(int indent = 0) const override;
    };

    class InsertStatement : public Statement
    {
    public:
        std::string table_name;
        std::vector<std::string> columns;
        std::vector<std::vector<std::shared_ptr<Expression>>> values;

        InsertStatement(std::string table_name_,
                        std::vector<std::string> columns_,
                        std::vector<std::vector<std::shared_ptr<Expression>>> values_)
            : table_name(std::move(table_name_)),
              columns(std::move(columns_)),
              values(std::move(values_)) {}

        std::string toString(int indent = 0) const override;
    };

    class UpdateStatement : public Statement
    {
    public:
        std::string table_name;
        std::vector<std::pair<std::string, std::shared_ptr<Expression>>> assignments;
        std::shared_ptr<Expression> where_clause;

        UpdateStatement(std::string table_name_,
                        std::vector<std::pair<std::string, std::shared_ptr<Expression>>> assignments_,
                        std::shared_ptr<Expression> where_clause_ = nullptr)
            : table_name(std::move(table_name_)),
              assignments(std::move(assignments_)),
              where_clause(std::move(where_clause_)) {}

        std::string toString(int indent = 0) const override;
    };

    class DeleteStatement : public Statement
    {
    public:
        std::string table_name;
        std::shared_ptr<Expression> where_clause;

        DeleteStatement(std::string table_name_, std::shared_ptr<Expression> where_clause_ = nullptr)
            : table_name(std::move(table_name_)), where_clause(std::move(where_clause_)) {}

        std::string toString(int indent = 0) const override;
    };

    class CreateTableStatement : public Statement
    {
    public:
        std::string table_name;
        std::vector<std::shared_ptr<ASTNode>> columns;

        CreateTableStatement(std::string table_name_, std::vector<std::shared_ptr<ASTNode>> columns_)
            : table_name(std::move(table_name_)), columns(std::move(columns_)) {}

        std::string toString(int indent = 0) const override;
    };

    class DropTableStatement : public Statement
    {
    public:
        std::string table_name;

        explicit DropTableStatement(std::string table_name_)
            : table_name(std::move(table_name_)) {}

        std::string toString(int indent = 0) const override;
    };

    // FROM子句
    class FromClause : public ASTNode
    {
    public:
        std::string table_name;
        std::string alias;

        FromClause(std::string table_name_, std::string alias_ = "")
            : table_name(std::move(table_name_)), alias(std::move(alias_)) {}

        std::string toString(int indent = 0) const override;
    };

    // 列定义
    class ColumnDefinition : public ASTNode
    {
    public:
        std::string column_name;
        std::string data_type;
        std::vector<std::string> constraints;

        ColumnDefinition(std::string column_name_, std::string data_type_,
                         std::vector<std::string> constraints_ = {})
            : column_name(std::move(column_name_)),
              data_type(std::move(data_type_)),
              constraints(std::move(constraints_)) {}

        std::string toString(int indent = 0) const override;
    };

    // 表达式类型
    class BinaryExpression : public Expression
    {
    public:
        std::shared_ptr<Expression> left;
        std::string operator_;
        std::shared_ptr<Expression> right;

        BinaryExpression(std::shared_ptr<Expression> left_,
                         std::string operator_,
                         std::shared_ptr<Expression> right_)
            : left(std::move(left_)), operator_(std::move(operator_)), right(std::move(right_)) {}

        std::string toString(int indent = 0) const override;
    };

    class UnaryExpression : public Expression
    {
    public:
        std::string operator_;
        std::shared_ptr<Expression> operand;

        UnaryExpression(std::string operator_, std::shared_ptr<Expression> operand_)
            : operator_(std::move(operator_)), operand(std::move(operand_)) {}

        std::string toString(int indent = 0) const override;
    };

    class ColumnReference : public Expression
    {
    public:
        std::string column_name;
        std::string table_name;

        ColumnReference(std::string column_name_, std::string table_name_ = "")
            : column_name(std::move(column_name_)), table_name(std::move(table_name_)) {}

        std::string toString(int indent = 0) const override;
    };

    class Literal : public Expression
    {
    public:
        std::variant<std::string, int, double> value;
        std::string value_type;

        Literal(std::variant<std::string, int, double> value_, std::string value_type_)
            : value(std::move(value_)), value_type(std::move(value_type_)) {}

        std::string toString(int indent = 0) const override;
    };

    class FunctionCall : public Expression
    {
    public:
        std::string function_name;
        std::vector<std::shared_ptr<Expression>> arguments;

        FunctionCall(std::string function_name_, std::vector<std::shared_ptr<Expression>> arguments_)
            : function_name(std::move(function_name_)), arguments(std::move(arguments_)) {}

        std::string toString(int indent = 0) const override;
    };

} // namespace MiniDB
